#!/bin/bash

# ABACUS势能面扫描脚本 - Be原子薄膜生长
# 菱形柱状区域扫描，同步执行确保任务顺序完成
# 每个扫描点创建独立文件夹保存结果

# 设置ABACUS路径和运行参数
ABACUS_CMD="/usr/local/bin/abacus"
CORE=10

# 菱形底面顶点坐标（埃）
A=(0 3.21)
B=(2.77994 4.815)
C=(2.77994 1.605)
D=(5.55988 3.21)

# 扫描参数（埃）
Z_START=4.2
Z_END=10.2
STEP=0.4

# 创建主结果文件
RESULT_FILE="potential_energy_surface.dat"
echo "# X Y Z Energy Force_X Force_Y Force_Z" > $RESULT_FILE

# 计算扫描范围
min_x=0
max_x=5.55988
min_y=1.605
max_y=4.815

echo "扫描范围: X[$min_x, $max_x], Y[$min_y, $max_y], Z[$Z_START, $Z_END]"
echo "扫描步长: $STEP 埃"

# 创建主结果目录
mkdir -p scan_results

# 实际扫描
total_points=0
completed_points=0

# 首先计算总点数
echo "计算总扫描点数..."
for x in $(seq $min_x $STEP $max_x); do
    for y in $(seq $min_y $STEP $max_y); do
        # 简化的菱形判断：使用线性不等式
        # 菱形可以看作两个等腰三角形组成
        # 上半部分：y <= -x + 7.02594 且 y <= x + 3.21
        # 下半部分：y >= x - 1.39594 且 y >= -x + 3.21
        
        # 转换为bc可计算的表达式
        cond1=$(echo "$y <= -$x + 7.02594" | bc -l)
        cond2=$(echo "$y <= $x + 3.21" | bc -l)
        cond3=$(echo "$y >= $x - 1.39594" | bc -l)
        cond4=$(echo "$y >= -$x + 3.21" | bc -l)
        
        # 如果所有条件都满足，点在菱形内
        if [ $cond1 -eq 1 ] && [ $cond2 -eq 1 ] && [ $cond3 -eq 1 ] && [ $cond4 -eq 1 ]; then
            for z in $(seq $Z_START $STEP $Z_END); do
                total_points=$((total_points + 1))
            done
        fi
    done
done

echo "总扫描点数: $total_points"
echo "开始扫描..."

# 实际扫描循环
for x in $(seq $min_x $STEP $max_x); do
    for y in $(seq $min_y $STEP $max_y); do
        # 简化的菱形判断：使用线性不等式
        cond1=$(echo "$y <= -$x + 7.02594" | bc -l)
        cond2=$(echo "$y <= $x + 3.21" | bc -l)
        cond3=$(echo "$y >= $x - 1.39594" | bc -l)
        cond4=$(echo "$y >= -$x + 3.21" | bc -l)
        
        # 如果所有条件都满足，点在菱形内
        if [ $cond1 -eq 1 ] && [ $cond2 -eq 1 ] && [ $cond3 -eq 1 ] && [ $cond4 -eq 1 ]; then
            for z in $(seq $Z_START $STEP $Z_END); do
                
                completed_points=$((completed_points + 1))
                progress=$(echo "scale=2; $completed_points * 100 / $total_points" | bc -l)
                echo "进度: $progress% - 扫描点 ($(printf "%.2f" $x), $(printf "%.2f" $y), $(printf "%.2f" $z))"
                
                # 创建扫描点目录
                dir_name=$(printf "scan_%.2f_%.2f_%.2f" $x $y $z)
                mkdir -p "scan_results/$dir_name"
                
                # 进入扫描点目录
                cd "scan_results/$dir_name"
                
                # 转换为Direct坐标
                X_dir=$(echo "scale=10; $x / 1.889726 / 11.1198" | bc -l)
                Y_dir=$(echo "scale=10; $y / 1.889726 / 9.63" | bc -l)
                Z_dir=$(echo "scale=10; $z / 1.889726 / 20.0" | bc -l)
                
                # 创建STRU文件
                cat > STRU << EOF
ATOMIC_SPECIES
Be    9.012  Be_ONCV_PBE-1.0.upf

NUMERICAL_ORBITAL
Be_gga_10au_100Ry_6s2p.orb

LATTICE_CONSTANT
1.889726

LATTICE_VECTORS
11.11980000000  0.000000000000  0.000000000000
0.000000000000  9.630000000000  0.000000000000
0.000000000000  0.000000000000  20.00000000000

ATOMIC_POSITIONS
Direct

Be
0.000
26
0.250000000000  0.166667000000  0.000000000000  1  1  1  mag  0.0
0.000000000000  0.666667000000  0.000000000000  1  1  1  mag  0.0
0.000000000000  0.000000000000  0.000000000000  1  1  1  mag  0.0
0.250000000000  0.500000000000  0.000000000000  1  1  1  mag  0.0
0.000000000000  0.333333000000  0.000000000000  1  1  1  mag  0.0
0.250000000000  0.833333000000  0.000000000000  1  1  1  mag  0.0
0.333333000000  0.000000000000  0.130325000000  1  1  1  mag  0.0
0.083333000000  0.500000000000  0.130325000000  1  1  1  mag  0.0
0.083333000000  0.166667000000  0.130325000000  1  1  1  mag  0.0
0.333333000000  0.666667000000  0.130325000000  1  1  1  mag  0.0
0.333333000000  0.333333000000  0.130325000000  1  1  1  mag  0.0
0.083333000000  0.833333000000  0.130325000000  1  1  1  mag  0.0
0.750000000000  0.166667000000  0.000000000000  1  1  1  mag  0.0
0.500000000000  0.666667000000  0.000000000000  1  1  1  mag  0.0
0.500000000000  0.000000000000  0.000000000000  1  1  1  mag  0.0
0.750000000000  0.500000000000  0.000000000000  1  1  1  mag  0.0
0.500000000000  0.333333000000  0.000000000000  1  1  1  mag  0.0
0.750000000000  0.833333000000  0.000000000000  1  1  1  mag  0.0
0.833333000000  0.000000000000  0.130325000000  1  1  1  mag  0.0
0.583333000000  0.500000000000  0.130325000000  1  1  1  mag  0.0
0.583333000000  0.166667000000  0.130325000000  1  1  1  mag  0.0
0.833333000000  0.666667000000  0.130325000000  1  1  1  mag  0.0
0.833333000000  0.333333000000  0.130325000000  1  1  1  mag  0.0
0.583333000000  0.833333000000  0.130325000000  1  1  1  mag  0.0
0.000000000000  0.333333000000  0.260650000000  1  1  1  mag  0.0
$(printf "%.10f" $X_dir)  $(printf "%.10f" $Y_dir)  $(printf "%.10f" $Z_dir)  1  1  1  mag  0.0
EOF
                
                # 复制INPUT文件到当前目录
                cp ../../INPUT_PARAMETERS INPUT
                
                # 后台运行ABACUS计算并等待完成
                echo "正在计算点 ($(printf "%.2f" $x), $(printf "%.2f" $y), $(printf "%.2f" $z))..."
                OMP_NUM_THREADS=1 mpirun -n $CORE $ABACUS_CMD > out.log 2> err.log &
                pid=$!
                wait $pid
                
                # 检查计算是否成功完成
                if [ $? -eq 0 ] && [ -f "OUT.ABACUS/running_scf.log" ]; then
                    # 提取能量
                    ENERGY=$(grep "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log | tail -1 | awk '{print $2}')
                    
                    # 从running_scf.log中提取力信息
                    # 查找包含"TOTAL-FORCE"的部分，然后提取最后一个原子的力
                    if grep -q "TOTAL-FORCE" OUT.ABACUS/running_scf.log; then
                        # 获取力的部分，并提取最后一个原子的力（第26个原子）
                        # 先找到最后一个TOTAL-FORCE部分
                        TOTAL_FORCE_SECTION=$(grep -n "TOTAL-FORCE" OUT.ABACUS/running_scf.log | tail -1 | cut -d: -f1)
                        if [ ! -z "$TOTAL_FORCE_SECTION" ]; then
                            # 跳过表头行，提取第26个原子的力
                            FORCE_LINE=$(tail -n +$((TOTAL_FORCE_SECTION + 3)) OUT.ABACUS/running_scf.log | head -n 26 | tail -n 1)
                            FORCE_X=$(echo $FORCE_LINE | awk '{print $2}')
                            FORCE_Y=$(echo $FORCE_LINE | awk '{print $3}')
                            FORCE_Z=$(echo $FORCE_LINE | awk '{print $4}')
                        else
                            FORCE_X="N/A"
                            FORCE_Y="N/A"
                            FORCE_Z="N/A"
                        fi
                    else
                        FORCE_X="N/A"
                        FORCE_Y="N/A"
                        FORCE_Z="N/A"
                    fi
                    
                    # 写入主结果文件（回到主目录）
                    cd ../..
                    echo "$x $y $z $ENERGY $FORCE_X $FORCE_Y $FORCE_Z" >> $RESULT_FILE
                    echo "计算完成: 能量 = $ENERGY, 力 = ($FORCE_X, $FORCE_Y, $FORCE_Z)"
                    
                    # 再次进入扫描点目录继续后续操作（如果需要）
                    cd "scan_results/$dir_name"
                else
                    # 回到主目录记录错误
                    cd ../..
                    echo "$x $y $z ERROR ERROR ERROR ERROR ERROR" >> $RESULT_FILE
                    echo "警告: 点 ($(printf "%.2f" $x), $(printf "%.2f" $y), $(printf "%.2f" $z)) 计算失败"
                    # 再次进入扫描点目录（如果需要清理或其他操作）
                    cd "scan_results/$dir_name"
                fi
                
                # 返回到主目录，准备下一个扫描点
                cd ../..
                
            done
        fi
    done
done

echo "扫描完成！结果保存在 $RESULT_FILE"
echo "所有计算结果保存在 scan_results/ 目录下的各个子目录中"
echo "总共完成 $completed_points/$total_points 个点的计算"